'''
ManiSkill 环境列表 - 完整版
'''

import gymnasium as gym
import mani_skill.envs
import numpy as np
import torch

def list_all_maniskill_envs():
    """列出所有可用的 ManiSkill 环境"""
    print("=" * 80)
    print("所有可用的 ManiSkill 环境:")
    print("=" * 80)
    
    # 获取所有注册的环境
    all_envs = list(gym.envs.registry.keys())
    
    # 分类收集 ManiSkill 环境
    maniskill_envs = {
        'Manipulation Tasks': [],
        'Scene Manipulation': [],
        'Mobile Manipulation': [],
        'Dexterous Manipulation': [],
        'Classic Control': [],
        'Other Tasks': []
    }
    
    # 关键词分类
    categories = {
        'Manipulation Tasks': ['pick', 'place', 'push', 'pull', 'stack', 'lift', 'cube', 'peg', 'door'],
        'Scene Manipulation': ['scene', 'tidy', 'groceries', 'table', 'replica'],
        'Mobile Manipulation': ['mobile', 'navigation', 'nav'],
        'Dexterous Manipulation': ['dexterous', 'hand', 'finger', 'grasp'],
        'Classic Control': ['cartpole', 'pendulum', 'balance'],
    }
    
    # 收集所有可能的 ManiSkill 环境
    potential_ms_envs = []
    for env_id in all_envs:
        # 多种匹配条件
        if any(keyword in env_id.lower() for keyword in [
            'ms-', 'mani', 'skill', 'sapien', 'cube', 'pick', 'push', 'stack', 
            'lift', 'peg', 'door', 'drawer', 'scene', 'replica', 'architec'
        ]):
            potential_ms_envs.append(env_id)
    
    # 分类环境
    for env_id in potential_ms_envs:
        categorized = False
        for category, keywords in categories.items():
            if any(keyword in env_id.lower() for keyword in keywords):
                maniskill_envs[category].append(env_id)
                categorized = True
                break
        
        if not categorized:
            maniskill_envs['Other Tasks'].append(env_id)
    
    # 打印分类结果
    total_count = 0
    for category, envs in maniskill_envs.items():
        if envs:
            print(f"\n📂 {category} ({len(envs)} 个环境):")
            print("-" * 60)
            for i, env_id in enumerate(sorted(envs), 1):
                print(f"  {i:2d}. {env_id}")
            total_count += len(envs)
    
    print(f"\n📊 总计: {total_count} 个 ManiSkill 相关环境")
    return potential_ms_envs

def get_env_details(env_id):
    """获取环境的详细信息"""
    try:
        # 尝试创建环境获取详细信息
        env = gym.make(env_id)
        
        details = {
            'id': env_id,
            'observation_space': str(env.observation_space),
            'action_space': str(env.action_space),
            'max_episode_steps': getattr(env.spec, 'max_episode_steps', 'Unknown'),
            'reward_threshold': getattr(env.spec, 'reward_threshold', 'Unknown'),
            'entry_point': getattr(env.spec, 'entry_point', 'Unknown'),
        }
        
        env.close()
        return details
    except Exception as e:
        return {'id': env_id, 'error': str(e)}

def analyze_environments(env_list, max_analyze=5):
    """分析前几个环境的详细信息"""
    print(f"\n🔍 详细分析前 {max_analyze} 个环境:")
    print("=" * 80)
    
    for i, env_id in enumerate(env_list[:max_analyze]):
        print(f"\n--- 环境 {i+1}: {env_id} ---")
        details = get_env_details(env_id)
        
        if 'error' in details:
            print(f"❌ 无法创建环境: {details['error']}")
        else:
            print(f"📏 观察空间: {details['observation_space']}")
            print(f"🎮 动作空间: {details['action_space']}")
            print(f"⏱️  最大步数: {details['max_episode_steps']}")
            print(f"🎯 奖励阈值: {details['reward_threshold']}")

def test_popular_envs():
    """测试一些热门的 ManiSkill 环境"""
    popular_envs = [
        "PushCube-v1",
        "PickCube-v1", 
        "StackCube-v1",
        "LiftCube-v1",
        "PegInsertionSide-v1",
        "OpenCabinetDoor-v1",
        "OpenCabinetDrawer-v1"
    ]
    
    print(f"\n🧪 测试热门环境:")
    print("=" * 80)
    
    working_envs = []
    
    for env_id in popular_envs:
        try:
            print(f"\n测试 {env_id}...", end=" ")
            env = gym.make(env_id, render_mode="rgb_array")
            
            # 简单测试
            obs, info = env.reset()
            action = env.action_space.sample()
            obs, reward, terminated, truncated, info = env.step(action)
            
            env.close()
            print("✅ 可用")
            working_envs.append(env_id)
            
        except Exception as e:
            print(f"❌ 不可用 ({str(e)[:50]}...)")
    
    print(f"\n✅ 可用的环境 ({len(working_envs)}/{len(popular_envs)}):")
    for env in working_envs:
        print(f"  • {env}")
    
    return working_envs

def search_environments(keyword):
    """根据关键词搜索环境"""
    all_envs = list(gym.envs.registry.keys())
    matched_envs = [env for env in all_envs if keyword.lower() in env.lower()]
    
    print(f"\n🔍 搜索关键词 '{keyword}' 的结果:")
    print("-" * 40)
    if matched_envs:
        for i, env in enumerate(matched_envs, 1):
            print(f"  {i:2d}. {env}")
        print(f"\n找到 {len(matched_envs)} 个匹配环境")
    else:
        print("  没有找到匹配的环境")
    
    return matched_envs

# 主程序
if __name__ == "__main__":
    print("🤖 ManiSkill 环境分析工具")
    
    # 1. 列出所有 ManiSkill 环境
    all_ms_envs = list_all_maniskill_envs()
    
    # 2. 详细分析前几个环境
    if all_ms_envs:
        analyze_environments(all_ms_envs, max_analyze=3)
    
    # 3. 测试热门环境
    working_envs = test_popular_envs()
    
    # 4. 提供搜索功能
    print(f"\n🔍 环境搜索示例:")
    keywords = ['cube', 'door', 'drawer', 'peg', 'scene']
    for keyword in keywords:
        matched = search_environments(keyword)
        if len(matched) > 0:
            break  # 只显示第一个有结果的搜索
    
    # 5. 推荐学习路径
    print(f"\n🎓 推荐的学习路径:")
    print("=" * 80)
    print("1. 初学者: PushCube-v1 (简单的推动任务)")
    print("2. 进阶: PickCube-v1 (抓取任务)")
    print("3. 高级: StackCube-v1 (堆叠任务)")
    print("4. 专家: 场景操作任务 (SceneManipulation 系列)")
    
    print(f"\n📚 更多信息:")
    print("- 官方文档: https://maniskill.readthedocs.io/")
    print("- GitHub: https://github.com/haosulab/ManiSkill")
    print("- 论文: https://arxiv.org/abs/2107.12564")